#!/usr/bin/env python
# coding=utf-8
'''
作者：林金城
邮箱：jinchengll@qq.com
描述：本脚本用来评估视频描述等任务的结果质量，评价指标有：
        - CIDEr
        - Bleu_4
        - Bleu_3
        - Bleu_2
        - Bleu_1
        - ROUGE_L
        - METEOR
'''
import argparse
import getopt
import hashlib
import io
import json
import os
import pylab
import sys
import re

sys.path.append('./coco-caption/')
from pycocotools.coco import COCO
from pycocoevalcap.eval import COCOEvalCap
from run_evaluations import CocoResFormat

class CocoResFormat:
  def __init__(self):
    self.res = []
    self.caption_dict = {}

  def read_multiple_files(self, filelist, hash_img_name):
    for filename in filelist:
      print 'In file %s\n' % filename
      self.read_file(filename, hash_img_name)

  def read_file(self, filename, hash_img_name):
    count = 0
    with open(filename,'r') as opfd:
      for line in opfd:
        count +=1
        id_sent = line.strip().split('\t')
        if len(id_sent)>2:
          id_sent = id_sent[-2:]
        assert len(id_sent) == 2
        sent = id_sent[1].decode('ascii', 'ignore')

        if hash_img_name:
          img_id = int(int(hashlib.sha256(id_sent[0]).hexdigest(),
                           16) % sys.maxint)
        else:  
          img = id_sent[0].split('_')[-1].split('.')[0]
          img_id = int(img)
        imgid_sent = {}
        
        if img_id in self.caption_dict:
          assert self.caption_dict[img_id] == sent
        else:
          self.caption_dict[img_id] = sent
          imgid_sent['image_id'] = img_id
          imgid_sent['caption'] = sent
          self.res.append(imgid_sent)
        if count%1000 == 0:
          print 'Processed %d ...' % count

  def dump_json(self, outfile):
    res = self.res
    with io.open(outfile, 'w', encoding='utf-8') as fd:
      fd.write(unicode(json.dumps(res,
         ensure_ascii=False,sort_keys=True,indent=2,separators=(',', ': '))))

caption_eval_path = './'
caption_output_path = './caption_output'
result_file = 'result.txt'
references_json_file = 'data/references.json'

def main():
  HASH_IMG_NAME = True
  pylab.rcParams['figure.figsize'] = (10.0, 8.0)
  json.encoder.FLOAT_REPR = lambda o: format(o, '.3f')

  # 获得output文件列表
  capitons_name = os.listdir(caption_output_path)
  capitons_name.sort(key=lambda x: int(re.findall('^\\d+', x)[0]))
  # 如果不存在的话，先生成references.json文件
  if not os.path.exists(references_json_file):
    cjr.out_api('data/references.txt', references_json_file)
  
  f_result = open(os.path.join(caption_eval_path, result_file), 'w')
  # 记录Meteor的最大值
  maxScore = {}

  coco = COCO(references_json_file)
  # 遍历计算每个描述文件的结果
  for caption in capitons_name:
    epcho, name = caption[:caption.index('p')], caption[caption.index('p'):]

    print '\n\n.....' + str(epcho) + ' is calculate.............'

    # 计算结果分数
    prediction_file = os.path.join(caption_output_path, caption)
    json_predictions_file = 'data/{0}.json'.format(prediction_file)

    crf = CocoResFormat()
    crf.read_file(prediction_file, HASH_IMG_NAME)
    crf.dump_json(json_predictions_file)
   
    cocoRes = coco.loadRes(json_predictions_file)

    cocoEval = COCOEvalCap(coco, cocoRes)
  
    # 计算指标
    cocoEval.evaluate()
  
    # 输出和保存结果
    f_result.write('\n'+'-'*10+str(epcho)+' epcho'+'-'*10+'\n')
    for metric, score in cocoEval.eval.items():
      print '%s: %.3f'%(metric, score)
      f_result.write('%s: %.3f\n'%(metric, score))
      # 记录最大值
      if not metric in maxScore or maxScore[metric] < score:
          maxScore[metric] = score
    #删除生成的prediction的json文件
    os.remove(json_predictions_file)

    print '............compled.....................'
  # 输出并保存最大值
  print '............MAXMAX.....................'
  f_result.write('-'*10+'MAXMAX'+'-'*10+'\n')
  for key in maxScore.keys():
      print '%s: %.3f'%(key, maxScore[key])
      f_result.write('%s: %.3f\n'%(key, maxScore[key]))
  f_result.close()

if __name__ == "__main__":
    main()
